""" Compare two structure files (POSCAR, CONTCAR, CHGCAR, etc.), and output
whether the atoms overlap. Usage
diffpos <file1> <file2>
The outputs are:
index label distance distance_x distance_y distance_z
The distance_{x,y,z} are in cartesian coordinates."""

import numpy as np
import sys

def frac2cart(b, f):
    """ transforms fractional coordinate to cartesian coordinate."""
    return np.array([
        b[0][0]*f[0] + b[0][1]*f[1] + b[0][2]*f[2],
        b[1][0]*f[0] + b[1][1]*f[1] + b[1][2]*f[2],
        b[2][0]*f[0] + b[2][1]*f[1] + b[2][2]*f[2]])

def read_to_cartesian(lines):
    """ read from POSCAR files, and convert to cartesian if not already."""
    # Line 0 of POSCAR file is a comment. The useful data are from Line 1.
    common_scalar = float(lines[1])
    basis = np.array([
        [float(x) for x in lines[2].split()],
        [float(x) for x in lines[3].split()],
        [float(x) for x in lines[4].split()]])
    basis = basis*common_scalar

    # Line 5 may be a list of element labels
    atom_species = lines[5].split()
    if atom_species[0].isdigit():
        # the old format
        number_atoms_subtype = [int(x) for x in atom_species]
        # I don't know species. Use XX
        atom_species = ["XX"] * len(number_atoms_subtype)
        line_next = 6
    else:
        # the new format
        number_atoms_subtype = [int(x) for x in lines[6].split()]
        line_next = 7

    # total number of atoms
    number_atoms_total = np.sum(number_atoms_subtype)

    if 'S'==lines[line_next][0] or 's'==lines[line_next][0] :
        # selective dynamics
        line_next += 1

    if lines[line_next][0] in "dD":
        is_fractional = True
    else:
        is_fractional = False
    line_next += 1

    spec = []
    coord = []
    for i, n in enumerate(number_atoms_subtype):
        for j in range(n):
            spec.append(atom_species[i])
            coord.append([float(x) for x in lines[line_next].split()[0:3]])
            line_next += 1
    if is_fractional:
        coord_c = []
        for x in coord:
            coord_c.append(frac2cart(basis, x))
        return [spec, np.array(coord_c)]
    else:
        return [spec, np.array(coord)]            


def dist3(v1, v2):
    """ calculate distance between two coordinates."""
    d = v1 - v2   # v1 and v2 are np.array
    return np.array([np.sqrt(np.dot(d,d)), d[0], d[1], d[2]])


if len(sys.argv) != 3:
    print("Usage: python diffpos.py POSCAR1 POSCAR2")
    sys.exit(1)

POS1 = open(sys.argv[1], "r")
all_lines1 = POS1.readlines()
POS1.close()

POS2 = open(sys.argv[2], "r")
all_lines2 = POS2.readlines()
POS2.close()

spec1, coord1 = read_to_cartesian(all_lines1)
spec2, coord2 = read_to_cartesian(all_lines2)

# The structures should be similar: they should have same number of atoms
# otherwise comparison is not meaningful at all
if len(spec1) != len(spec2):
    print("The two structures have different numbers of atoms. I'll stop.")
    sys.exit(2)

for n, label in enumerate(spec1):
    d = dist3(coord1[n,:], coord2[n,:])
    sys.stdout.write("{0:d} {1:s}: {2:9.6f} {3:7.3f} {4:7.3f} {5:7.3f}\n".format(
        n, label, d[0], d[1], d[2], d[3]))


